import torch
import torch.nn as nn
from torch.nn import init


class SupervisedGraphSage(nn.Module):

    def __init__(self, num_classes, enc):
        super(SupervisedGraphSage, self).__init__()
        # 这里面赋值为enc2(经过两层GNN)
        self.enc = enc
        self.xent = nn.CrossEntropyLoss()
        # 全连接参数矩阵，映射到labels num_classes维度做分类
        self.weight = nn.Parameter(torch.FloatTensor(num_classes, enc.embed_dim))
        init.xavier_uniform(self.weight)

    def forward(self, nodes):
        # embeds实际是我们两层GNN后的输出nodes embedding
        embeds = self.enc(nodes)
        # 最后将nodes * hidden size 映射到 nodes * num_classes(= 7)之后做softmax计算cross entropy
        scores = self.weight.mm(embeds)
        return scores.t()

    def loss(self, nodes, labels):
        # 钱箱传播
        scores = self.forward(nodes)
        # 定义的cross entropy
        return self.xent(scores, labels.squeeze())
